#!/usr/bin/env python3
"""
P4Runtime源路由控制器实现
"""

import argparse
import grpc
import os
import sys
import time
from typing import Dict, List, Tuple, Any

# 导入P4Runtime库
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper


class SwitchConnectionConfig:
    """交换机连接配置"""
    SWITCH_LIST = [
        {'name': 's1', 'address': '127.0.0.1:50051', 'device_id': 0},
        {'name': 's2', 'address': '127.0.0.1:50052', 'device_id': 1},
        {'name': 's3', 'address': '127.0.0.1:50053', 'device_id': 2}
    ]


class RouteTableManager:
    """管理路由表项"""
    
    @staticmethod
    def add_table_entry(p4info_helper, switch, table_name: str, 
                         match_fields: Dict, action_name: str, 
                         action_params: Dict[str, Any]) -> None:
        """向交换机添加表项"""
        table_entry = p4info_helper.buildTableEntry(
            table_name=table_name,
            match_fields=match_fields,
            action_name=action_name,
            action_params=action_params
        )
        switch.WriteTableEntry(table_entry)
    
    @staticmethod
    def configure_mac_rewrite(p4info_helper, switch, 
                               dest_ip: str, prefix_len: int,
                               mac_addr: str) -> None:
        """配置MAC地址重写规则"""
        RouteTableManager.add_table_entry(
            p4info_helper=p4info_helper,
            switch=switch,
            table_name="MyIngress.ipv4_lpm",
            match_fields={"hdr.ipv4.dstAddr": [dest_ip, prefix_len]},
            action_name="MyIngress.rewriteMac",
            action_params={"dstAddr": mac_addr}
        )
    
    @staticmethod
    def configure_source_routes(p4info_helper, switch, 
                                  dest_ip: str, prefix_len: int,
                                  action_name: str, hop_params: Dict[str, Any]) -> None:
        """配置源路由表项"""
        RouteTableManager.add_table_entry(
            p4info_helper=p4info_helper,
            switch=switch,
            table_name="MyIngress.src_routing_publish",
            match_fields={"hdr.ipv4.dstAddr": [dest_ip, prefix_len]},
            action_name=action_name,
            action_params=hop_params
        )


class NetworkRoutingConfig:
    """网络路由配置"""
    
    # 主机配置
    HOST_SETTINGS = {
        'h1': {'ip': '10.0.1.1', 'mac': '08:00:00:00:01:11'},
        'h2': {'ip': '10.0.2.2', 'mac': '08:00:00:00:02:22'},
        'h3': {'ip': '10.0.3.3', 'mac': '08:00:00:00:03:33'}
    }
    
    # 交换机路由规则
    SWITCH_ROUTES = {
        's1': [
            # 1跳路由
            ('10.0.1.1', 32, 'MyIngress.add_head_1', {'bos1': 1, 'port1': 1}),
            # 2跳路由
            ('10.0.2.2', 32, 'MyIngress.add_head_2', {'bos1': 0, 'port1': 2, 'bos2': 1, 'port2': 1}),
            ('10.0.3.3', 32, 'MyIngress.add_head_2', {'bos1': 0, 'port1': 3, 'bos2': 1, 'port2': 1})
        ],
        's2': [
            # 1跳路由
            ('10.0.2.2', 32, 'MyIngress.add_head_1', {'bos1': 1, 'port1': 1}),
            # 2跳路由
            ('10.0.1.1', 32, 'MyIngress.add_head_2', {'bos1': 0, 'port1': 2, 'bos2': 1, 'port2': 1}),
            ('10.0.3.3', 32, 'MyIngress.add_head_2', {'bos1': 0, 'port1': 3, 'bos2': 1, 'port2': 1})
        ],
        's3': [
            # 1跳路由
            ('10.0.3.3', 32, 'MyIngress.add_head_1', {'bos1': 1, 'port1': 1}),
            # 2跳路由
            ('10.0.1.1', 32, 'MyIngress.add_head_2', {'bos1': 0, 'port1': 2, 'bos2': 1, 'port2': 1}),
            ('10.0.2.2', 32, 'MyIngress.add_head_2', {'bos1': 0, 'port1': 3, 'bos2': 1, 'port2': 1})
        ]
    }


class SwitchController:
    """交换机控制器"""
    
    def __init__(self, p4info_helper, bmv2_file: str):
        self.p4info_helper = p4info_helper
        self.bmv2_file = bmv2_file
        self.switches = {}
    
    def connect_switches(self) -> None:
        """连接并初始化所有交换机"""
        print("初始化交换机连接...")
        
        for switch_cfg in SwitchConnectionConfig.SWITCH_LIST:
            self._setup_switch(switch_cfg)
        
        # 建立主控制器连接并安装P4程序
        for switch_name, switch in self.switches.items():
            print(f"配置 {switch_name}...")
            switch.MasterArbitrationUpdate()
            switch.SetForwardingPipelineConfig(
                p4info=self.p4info_helper.p4info,
                bmv2_json_file_path=self.bmv2_file
            )
    
    def _setup_switch(self, config: Dict) -> None:
        """建立交换机连接"""
        proto_dump = f"logs/{config['name']}-p4runtime-requests.txt"
        self.switches[config['name']] = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name=config['name'],
            address=config['address'],
            device_id=config['device_id'],
            proto_dump_file=proto_dump
        )
    
    def setup_routing(self) -> None:
        """配置所有交换机的路由表"""
        print("配置路由表...")
        
        # 配置源路由
        self._setup_source_routing()
        
        # 配置MAC重写
        self._setup_mac_rewrite()
        
        print("路由配置完成。")
    
    def _setup_source_routing(self) -> None:
        """配置源路由表"""
        for switch_name, routes in NetworkRoutingConfig.SWITCH_ROUTES.items():
            switch = self.switches[switch_name]
            print(f"配置 {switch_name} 的源路由...")
            
            for dest_ip, prefix_len, action, params in routes:
                RouteTableManager.configure_source_routes(
                    p4info_helper=self.p4info_helper,
                    switch=switch,
                    dest_ip=dest_ip,
                    prefix_len=prefix_len,
                    action_name=action,
                    hop_params=params
                )
    
    def _setup_mac_rewrite(self) -> None:
        """配置MAC重写表"""
        print("配置MAC重写表...")
        
        mac_setup = [
            ('s1', '10.0.1.1', 32, NetworkRoutingConfig.HOST_SETTINGS['h1']['mac']),
            ('s2', '10.0.2.2', 32, NetworkRoutingConfig.HOST_SETTINGS['h2']['mac']),
            ('s3', '10.0.3.3', 32, NetworkRoutingConfig.HOST_SETTINGS['h3']['mac'])
        ]
        
        for switch_name, ip, prefix, mac in mac_setup:
            RouteTableManager.configure_mac_rewrite(
                p4info_helper=self.p4info_helper,
                switch=self.switches[switch_name],
                dest_ip=ip,
                prefix_len=prefix,
                mac_addr=mac
            )


def check_files(p4info_path: str, bmv2_json_path: str) -> None:
    """检查文件是否存在"""
    if not os.path.exists(p4info_path):
        raise FileNotFoundError(f"p4info文件不存在: {p4info_path}")
    if not os.path.exists(bmv2_json_path):
        raise FileNotFoundError(f"BMv2 JSON文件不存在: {bmv2_json_path}")


def show_routing_info() -> None:
    """显示路由配置摘要"""
    print("\n" + "="*50)
    print("源路由配置摘要")
    print("="*50)
    
    for host, cfg in NetworkRoutingConfig.HOST_SETTINGS.items():
        print(f"主机 {host}: IP={cfg['ip']}, MAC={cfg['mac']}")
    
    print("\n交换机路由规则:")
    for switch, rules in NetworkRoutingConfig.SWITCH_ROUTES.items():
        print(f"\n{switch}:")
        for ip, prefix, action, params in rules:
            hop_count = action.split('_')[-1]
            print(f"  {ip}/{prefix} -> {action} (参数: {params})")


def main(p4info_file: str, bmv2_file: str) -> None:
    """主控制器函数"""
    check_files(p4info_file, bmv2_file)
    
    # 初始化P4Runtime助手
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file)
    
    try:
        # 设置交换机控制器
        switch_mgr = SwitchController(p4info_helper, bmv2_file)
        switch_mgr.connect_switches()
        
        # 配置路由表
        switch_mgr.setup_routing()
        
        # 显示配置摘要
        show_routing_info()
        
        print("\n源路由控制器运行中。按Ctrl+C退出。")
        
        # 保持运行
        while True:
            time.sleep(1)
            
    except KeyboardInterrupt:
        print("\n关闭源路由控制器。")
    except grpc.RpcError as e:
        printGrpcError(e)
    except Exception as e:
        print(f"意外错误: {e}")
    finally:
        ShutdownAllSwitchConnections()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='P4Runtime源路由控制器',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    
    parser.add_argument(
        '--p4info',
        help='p4c生成的p4info文本格式',
        type=str,
        default='./build/source_routing_v2.p4.p4info.txtpb'
    )
    
    parser.add_argument(
        '--bmv2-json',
        help='p4c生成的BMv2 JSON文件',
        type=str,
        default='./build/source_routing_v2.json'
    )
    
    args = parser.parse_args()
    
    try:
        main(args.p4info, args.bmv2_json)
    except FileNotFoundError as e:
        print(f"错误: {e}")
        print("是否运行了'make'来构建所需文件?")
        sys.exit(1)